import cv2
import numpy as np
class ResizeImage():
    """
    图片resize, 可以用于避免输入图像过大导致爆显存等问题

    args:
        keys: 指定了从data中读取数据的关键字和保存数据的关键字
            格式: keys={"in":"your_in_keyword",
                        "out":"your_save_keyword"
                        }
    """
    def __init__(self, keys,max_size):
        self.keys = keys
        self.max_size = max_size

    def __call__(self, data:dict):
        img_data_list = data[self.keys["in"]]
        out_data_list = []
        for img_data in img_data_list:
            shape = img_data.shape
            scale = max(max(shape)/self.max_size,1)
            if scale >1:
                new_shape = [int(shape[1]/scale)//2*2, int(shape[0]/scale)//2*2]
                new_imgdata = cv2.resize(img_data, new_shape)
            else:
                new_imgdata = img_data
            out_data_list.append(new_imgdata)
        data[self.keys["out"]] = out_data_list
        return data